#ifndef DATASET_H
#define DATASET_H

#include <iostream>
#include <vector>
#include <string>
#include <filesystem>
#include <algorithm>
#include <fstream>
#include "type_alias.h"

#include <opencv2/opencv.hpp>
#include <optional>

#define CUSTOM_ASSERT(condition, message) \
    do { \
        if (!(condition)) { \
            std::cerr << "Assertion failed: " << message << " | " \
                      << "File: " << __FILE__ << " | " \
                      << "Line: " << __LINE__ << " | " \
                      << "Function: " << __func__ << "\n"; \
            std::abort(); \
        } \
    } while (0)

struct CUSTOM_INFO {
    const std::string alphaNum = " 0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ+&~!@#";
    int max_height = 100;    // train image.height
    int max_width = 200;     // train image.width
};




using Example = std::pair<Mat3d, std::vector<size_t>>;
using VecMat = std::vector<cv::Mat>;
using VecInt = std::vector<int>;
using ImgLable = std::pair<cv::Mat, std::vector<int>>;



class OCRDataset{
public:
    explicit OCRDataset(const std::string &root_dir) {
        // 遍历目录获取所有PNG文件
        for (const auto &entry: std::filesystem::directory_iterator(root_dir)) {
            if (entry.path().extension() != ".png") { continue; }

            std::string filename = entry.path().stem().string();

            // 解析文件名中的标签
            size_t last_underscore = filename.find_last_of('_');
            if (last_underscore == std::string::npos) { continue; }
            std::string label_str = filename.substr(0, last_underscore);
            if (label_str.empty()) { continue; }

            // 转换字符到索引
            std::vector<int> label_indices;
            try {
                label_indices = string_to_indices(label_str);
                samples_.emplace_back(entry.path().string(), label_indices);
            } catch (const std::exception &e) {
                std::cerr << "Error parsing label string: " << label_str << " - " << e.what() << std::endl;
                continue;
            }
        }

        unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
        std::default_random_engine rng(seed);
        std::shuffle(samples_.begin(), samples_.end(), rng);
    }

    std::pair<cv::Mat, Mat1i> get(size_t index) {
        std::string img_path = samples_[index].first;
        cv::Mat image = cv::imread(img_path, cv::IMREAD_COLOR);
        CUSTOM_ASSERT(image.rows == cust_info.max_height && image.cols == cust_info.max_width, "Image size is not correct");
        cv::cvtColor(image, image, cv::COLOR_BGR2RGB);

        // Convert to tensor   HWC -> CHW
        /*
         torch::Tensor tensor_image = torch::from_blob(
                image.data,
                {image.rows, image.cols, 3},
                torch::kByte
        ).permute({2, 0, 1})
                .to(torch::kFloat32)
                .div(255)
                .sub(0.5)
                .div(0.5);


        // 创建标签张量
        std::vector<int64_t> label(samples_[index].second.begin(), samples_[index].second.end());
        torch::Tensor tensor_label = torch::from_blob(
                label.data(), {static_cast<int64_t>(label.size())}, torch::kInt64).clone();
        */

        return {image, samples_[index].second};
    }

    [[nodiscard]] std::optional<size_t> size() const { return samples_.size();}

private:
    std::vector<std::pair<std::string, std::vector<int>>> samples_;
    CUSTOM_INFO cust_info;

    std::vector<int> string_to_indices(const std::string &str) {
        std::vector<int> indices;
        for (char c: str) {
            size_t pos = cust_info.alphaNum.find(c);
            if (pos == std::string::npos) {
                throw std::runtime_error("Invalid character in label: " + std::string(1, c));
            }
            indices.push_back(static_cast<int>(pos));
        }
        return indices;
    }

};



std::tuple<Mat4d, Mat2i, Mat1i> collate_fn_vec(std::vector<ImgLable>& batch) {
    std::vector<int> sequence_lengths;
    sequence_lengths.reserve(batch.size());

    int max_width = CUSTOM_INFO().max_width;
    int max_height = CUSTOM_INFO().max_height;

    for (const auto& example : batch) {
        sequence_lengths.push_back((int)(example.second.size()));
    }

    Mat4d img_tensor(batch.size(), Mat3d(3, Mat2d(max_height, Mat1d(max_width, 0.0))));
    Mat2i seq_tensor(batch.size(), std::vector<int>(*std::max_element(sequence_lengths.begin(), sequence_lengths.end()), 0));

    for (size_t idx = 0; idx < batch.size(); ++idx) {
        auto& image = batch[idx].first;
        auto& label = batch[idx].second;
        for (size_t i = 0; i < label.size(); ++i) {
            seq_tensor[idx][i] = label[i];
        }

        for (int row = 0; row < max_height; ++row) {
            auto* pdata = image.ptr<uchar>(row);
            for (int col = 0; col < max_width; ++col) {
                double pixel_value0 = static_cast<double>(*pdata++) / 255.0; // pixel -> 0-1
                double pixel_value1 = static_cast<double>(*pdata++) / 255.0; // pixel -> 0-1
                double pixel_value2 = static_cast<double>(*pdata++) / 255.0; // pixel -> 0-1

                pixel_value0 = (pixel_value0 - 0.5) / 0.5;
                pixel_value1 = (pixel_value1 - 0.5) / 0.5;
                pixel_value2 = (pixel_value2 - 0.5) / 0.5;

                img_tensor[idx][0][row][col] = pixel_value0;
                img_tensor[idx][1][row][col] = pixel_value1;
                img_tensor[idx][2][row][col] = pixel_value2;
            }
        }
    }

    return std::make_tuple(img_tensor, seq_tensor, sequence_lengths);
}



std::tuple<Tensor, Mat2i, Mat1i> collate_fn(std::vector<ImgLable>& batch) {
    std::vector<int> sequence_lengths;
    sequence_lengths.reserve(batch.size());

    int max_width = CUSTOM_INFO().max_width;
    int max_height = CUSTOM_INFO().max_height;

    for (const auto& example : batch) {
        sequence_lengths.push_back((int)(example.second.size()));
    }

    Tensor img_tensor(batch.size(), 3, max_height, max_width);
    Mat2i seq_tensor(batch.size(), std::vector<int>(*std::max_element(sequence_lengths.begin(), sequence_lengths.end()), 0));

    for (size_t idx = 0; idx < batch.size(); ++idx) {
        auto& image = batch[idx].first;
        auto& label = batch[idx].second;
        for (size_t i = 0; i < label.size(); ++i) {
            seq_tensor[idx][i] = label[i];
        }

        for (int row = 0; row < max_height; ++row) {
            auto* pdata = image.ptr<uchar>(row);
            for (int col = 0; col < max_width; ++col) {
                double pixel_value0 = static_cast<double>(*pdata++) / 255.0; // pixel -> 0-1
                double pixel_value1 = static_cast<double>(*pdata++) / 255.0; // pixel -> 0-1
                double pixel_value2 = static_cast<double>(*pdata++) / 255.0; // pixel -> 0-1

                pixel_value0 = (pixel_value0 - 0.5) / 0.5;
                pixel_value1 = (pixel_value1 - 0.5) / 0.5;
                pixel_value2 = (pixel_value2 - 0.5) / 0.5;

                img_tensor.at(idx, 0, row, col) = pixel_value0;
                img_tensor.at(idx, 1, row, col) = pixel_value1;
                img_tensor.at(idx, 2, row, col) = pixel_value2;
            }
        }
    }

    return std::make_tuple(img_tensor, seq_tensor, sequence_lengths);
}


#endif